package com.frost.case2;

import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;

import java.io.IOException;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;

public class DistributedLock {

    private final String connectString = "centos102:2181,centos103:2181,centos104:2181";
    private final int sessionTimeout = 2000;
    ZooKeeper zk;

    private CountDownLatch connectLatch = new CountDownLatch(1);
    private CountDownLatch waitLatch = new CountDownLatch(1);

    // 前一个节点
    private String waitPath;
    // 当前节点
    String currentMode;

    public DistributedLock() throws IOException, InterruptedException, KeeperException {
        // 1. 获取连接
         zk = new ZooKeeper(connectString, sessionTimeout, new Watcher() {
            @Override
            public void process(WatchedEvent event) {
                // connectLatch，如果连接上zk，可以释放
                if (event.getState() == Event.KeeperState.SyncConnected) {
                    connectLatch.countDown();
                }

                // waitLatch，需要释放
                if (event.getType() == Event.EventType.NodeDeleted && event.getPath().equals(waitPath)) {
                    waitLatch.countDown();
                }
            }
        });

         // 等待zk正常连接后往下走
         connectLatch.await();

        // 2. 判断根节点/lock是否存在
        Stat stat = zk.exists("/locks", false);

        if (stat == null) {
            // 创建根节点(永久节点)
            zk.create("/locks", "locks".getBytes(), ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
        }
    }

    // 对 zk 加锁
    public void zkLock() {
        // 创建对应的临时带序号的节点
        try {
            currentMode = zk.create("/locks/" + "seq-", null, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL);

            // 判断创建的节点是否是最小的序号节点，如果是，获取到锁；如果不是，监听前一个节点
            List<String> children = zk.getChildren("/locks", false);

            // 如果 children 只有一个节点，直接获取锁；如果有多个节点，需要判断，谁最小
            if (children.size() == 1) {
                return;
            }else {
                // 排序
                Collections.sort(children);

                // 获取节点名称 seq00000001
                String thisNode = currentMode.substring("/locks/".length());
                // 通过 seq00000001 获取该节点在children当中的位置
                int index = children.indexOf(thisNode);

                if (index == -1) {
                    System.out.println("数据异常");
                }else if (index == 0) {
                    // 该节点为第一个，获取锁直接返回
                    return;
                }else {
                    // 不是第一个，监听前一个节点
                    waitPath = "/locks/" + children.get(index - 1);
                    zk.getData(waitPath, true, null);

                    // 等待监听结束
                    waitLatch.await();
                    return;
                }
            }
        } catch (KeeperException e) {
            e.printStackTrace();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
    }

    // 解锁
    public void unZkLock() throws InterruptedException, KeeperException {
        // 删除节点
        zk.delete(currentMode, -1);
    }
}
